import numpy as np
from PIL import Image
from scipy import ndimage
from glob import glob
import os
from tqdm import tqdm
from keras.utils import to_categorical

def formatImage(img, format_size=(128, 128)):
    image = img
    fg_w, fg_h = image.size
    bg_w, bg_h = format_size
    
    if image.mode == 'L':
        out_image = np.zeros((bg_h, bg_w), dtype='uint8') + 255
    elif image.mode == 'RGB':
        out_image = np.zeros((bg_h, bg_w, 3), dtype='uint8') + 255
    elif image.mode == 'RGBA':
        out_image = np.zeros((bg_h, bg_w, 4), dtype='uint8') + 255
    else:
        image = image.convert('RGB')
        out_image = np.zeros((bg_h, bg_w, 3), dtype='uint8') + 255

    if fg_h == fg_w:
        image2 = image.resize(format_size, resample=Image.BOX)
        out_image = np.array(image2, dtype='uint8')
    elif fg_h > fg_w:
        h = bg_h
        w = int(h / fg_h * fg_w)
        image2 = image.resize((w, h), resample=Image.BOX)
        image2_arr = np.array(image2, dtype='uint8')
        out_image[:, (bg_w - w) // 2: (bg_w - w) // 2 + w] = image2_arr
    elif fg_w > fg_h:
        w = bg_w
        h = int(w / fg_w * fg_h)
        image2 = image.resize((w, h), resample=Image.BOX)
        image2_arr = np.array(image2, dtype='uint8')
        out_image[(bg_h - h) // 2: (bg_h - h) // 2 + h, :] = image2_arr
    return out_image

def loadDatasets(path, format_size=(128, 128)):
    paths = glob(os.path.join(path, '*'))
    datas = []
    targets = []
    for path in tqdm(paths, ncols=80):
        img = Image.open(path)
        if img.mode != 'L':
            img = img.convert('L')
        else:
            img = img
        img_arr = formatImage(img, format_size)

        for time in range(4):
            k = np.random.randint(0, 4)
            rotated = np.rot90(img_arr, k=k)
            datas.append(rotated)
            targets.append(k)
    return datas, targets

def datas2X(datas, format_size=(128, 128)):
    train_X = np.array(datas, dtype='float32')
    train_X = train_X.reshape(-1, format_size[0], format_size[1], 1)
    train_X = 1 - train_X / 255
    return train_X

def targets2Y(targets):
    train_Y = to_categorical(np.array(targets), num_classes=4)
    return train_Y
